﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Linq;
using System.Text;
using System.Windows.Forms;
using SharpRBM.Core;
using System.IO;
using SharpRBM.DemoTools;
using System.Diagnostics;

namespace SharpRBM.ImageningFaces
{
    public partial class MainForm : Form
    {

        public static string DeepBeliefNetworkFile = Variables.DataPath + "ImageningFaces.dbn";
        //public const string ImagePath = DataPath+@"SuperSimpleFeatures\img{1}.bmp";


        public MainForm()
        {
            InitializeComponent();
            LoadDeepBeliefNetwork();
        }

        public bool Stopping
        {
            get
            {
                return stopButton.Enabled == false;
            }
            set
            {
                stopButton.Enabled = !value;
                startButton.Enabled = !stopButton.Enabled;
            }
        }

        public DeepBeliefNetwork DeepBeliefNetwork { get; set; }

        private List<Layer> _inputLayers;
        private bool _unsavedChanges = false;

        private void LoadDeepBeliefNetwork()
        {
            _unsavedChanges = false;
            if (File.Exists(DeepBeliefNetworkFile))
            {
                AddLog("Loading DeepBeliefNetwork");
                DeepBeliefNetwork = SerializationHelper.LoadFromFile<DeepBeliefNetwork>(DeepBeliefNetworkFile);
            }
            else
            {
                AddLog("Creating new DeepBeliefNetwork");
                // First layer is image layer, the next two are hidden layer 1 and hidden layer 2
                DeepBeliefNetwork = new DeepBeliefNetwork(
                    LayerBitmapHelper.BitmapSize * LayerBitmapHelper.BitmapSize + 1,
                    Variables.SecondLayerSize,
                    Variables.ThirdLayerSize,
                    Variables.FourthLayerSize );
            }

            DocumentDeepBeliefNetwork();
        }

        private void DocumentDeepBeliefNetwork()
        {
            AddLog("Network contains {0} layers:", DeepBeliefNetwork.LayerSizes.Count);

            Helper.Loop(
                DeepBeliefNetwork.LayerSizes.Count,
                i =>
                {
                    string add = "";
                    if (i < DeepBeliefNetwork.LayerSizes.Count - 1)
                    {
                        InterLayerWeights interLayerWeights = DeepBeliefNetwork.LayerWeights[i];
                        add = string.Format(
                            ", TrainingSteps={2}, ReconstructionError={3:0}, DetectorError={4:0}, Lower={5}, Upper={6}",
                            i,
                            i + 1,
                            interLayerWeights.PerformedTrainingSteps,
                            interLayerWeights.LastTrainingError != null ? interLayerWeights.LastTrainingError.ReconstructionError : -1,
                            interLayerWeights.LastTrainingError != null ? interLayerWeights.LastTrainingError.FeatureDetectorError : -1,
                            interLayerWeights.LowerLayerSize,
                            interLayerWeights.UpperLayerSize
                            );
                    }
                    else
                    {
                        add = ", no training";
                    }

                    AddLog("  Layer[{0}]: Nodes={1}{2}", i, DeepBeliefNetwork.LayerSizes[i], add);
                });

            AddLog("");
        }

        private void AddLog(string message)
        {
            if (IsDisposed == false && logTextBox.IsDisposed == false)
            {
                logTextBox.AppendText(message + Environment.NewLine);
            }
        }

        private void AddLog(string message, params object[] args)
        {
            AddLog(string.Format(message, args));
        }

        private void Form1_Shown(object sender, EventArgs e)
        {
            if (VerifyThatDataZipHasBeenUncompressed() == false)
            {
                Close();
                return;
            }

            Stopping = true;
            ClearBoxes();
            LoadInputs();
        }

        private bool VerifyThatDataZipHasBeenUncompressed()
        {
            FileInfo fileInfo = new FileInfo( Variables.DataPath ); // get a FileInfo for the file
            while( Directory.Exists( Variables.VerifyPath ) == false )
            {
                DialogResult result =
                    MessageBox.Show(
                    this,
                    string.Format(
                        "Before you can run this program, you need to uncompress the {0}\\Data.Zip to {0}\\.\nPress OK when you've uncompresed the file and Cancel if you want to quit.",
                        fileInfo.Directory),
                    "SharpRBM",
                    MessageBoxButtons.OKCancel);

                if (result == DialogResult.Cancel)
                {
                    return false;
                }
            }
            return true;
        }

        private void ClearBoxes()
        {
            //inputBictureBox.Image = new Bitmap(inputBictureBox.Width, inputBictureBox.Height);
            reconstructedPictureBox.Image = new Bitmap(reconstructedPictureBox.Width, reconstructedPictureBox.Height);
        }

        private void ShowFeatureDetectors(DeepBeliefNetworkTrainer trainer)
        {
            if (drawFeatureDetectorsCheckBox.Checked == false)
            {
                return;
            }

            Layer upper = new Layer(trainer.TrainingWeights.UpperLayerSize);
            int shownCount = ShowBitmaps(
                featureDetectors,
                Enumerable.Range(0, upper.Size), i =>
                {
                    upper.SetValues(0);

                    if (enableBiasDuringFeatureDetectorRenderCheckBox.Checked)
                    {
                        upper.SetBias();
                    }
                    else
                    {
                        upper.Values[0] = 0;
                    }

                    upper.Values[i] = 1;

                    Layer reconstructed = trainer.ActivateDownToInput(upper);

                    if (increaseContrastCheckBox.Checked)
                    {
                        return LayerBitmapHelper.CreateContrastEnhancedBitmapFromLayer(reconstructed);
                    }
                    else
                    {
                        return LayerBitmapHelper.CreateBitmapFromLayer(reconstructed);
                    }
                });

            featureDetectorsGroupBox.Text = string.Format("Feature detectors (first {0} out of {1})", shownCount, upper.Size);
        }

        private void ShowInputs(List<Bitmap> bitmaps)
        {
            int shownCount = ShowBitmaps(inputBictureBox, bitmaps, bitmap => bitmap);
            inputsLabel.Text = string.Format("Inputs (first {0} our of {1})", shownCount, bitmaps.Count);
        }

        private static int ShowBitmaps<T>(PictureBox target, IEnumerable<T> items, Func<T, Bitmap> func)
        {
            if (target.Height == 0 || target.Width == 0)
            {
                return 0;
            }

            Bitmap bitmap = new Bitmap(target.Width, target.Height);
            int count = 0;
            using (Graphics graphics = Graphics.FromImage(bitmap))
            {
                int x = 0;
                int y = 0;
                int width = LayerBitmapHelper.BitmapSize + LayerBitmapHelper.BitmapMargin;
                foreach (T item in items)
                {
                    if (x + width > target.Width)
                    {
                        x = 0;
                        y += width;

                        if (y + width > target.Height)
                        {
                            break;
                        }
                    }

                    graphics.DrawImageUnscaled(func(item), new Point(x, y));
                    count++;

                    x += width;
                }

                target.Image = bitmap;
            }
            return count;
        }

        private void ShowReconstructed(DeepBeliefNetworkTrainer trainer)
        {
            if (drawReconstructionsCheckBox.Checked)
            {
                int shownCount = ShowBitmaps(
                    reconstructedPictureBox,
                    _inputLayers,
                    layer =>
                    {
                        Layer reconstruction = trainer.Reconstruct(layer, false);
                        return LayerBitmapHelper.CreateBitmapFromLayer(reconstruction);
                    });

                reconstructedLabel.Text = string.Format("Reconstructed (first {0} out of {1})", shownCount, _inputLayers.Count);
            }
        }

        public double GetErrorSum(InterLayerWeights weights, List<Layer> inputs, int upperLayerSize)
        {
            Layer testLower = new Layer(inputs[0].Size);
            Layer upper = new Layer(upperLayerSize);

            return
                inputs.Sum(layer => ComputeError(weights, layer, upper, testLower));
        }

        public double ComputeError(InterLayerWeights weights, Layer lowerLayer, Layer upperLayer, Layer testLower)
        {
            weights.ActivateLowerToUpper(lowerLayer, upperLayer);
            weights.ActivateUpperToLower(testLower, upperLayer);

            double errorSum = 0;
            Helper.Loop(lowerLayer.Size, i =>
                {
                    double error = lowerLayer.Values[i] - testLower.Values[i];
                    errorSum += error * error;
                });

            return errorSum;
        }

        private void LoadInputs()
        {
            // Input
            List<Bitmap> bitmaps = GetInputBitmaps( Variables.InputImageCount );
            ShowInputs(bitmaps);

            // Input layers...
            _inputLayers = GetInputLayers(bitmaps);
        }

        public List<Bitmap> GetInputBitmaps(int count)
        {
            List<Bitmap> bitmaps = new List<Bitmap>();
            for (int i = 0; i < count; i++)
            {
                for (int number = 0; number < 10; number++)
                {
                    string fileName = string.Format( Variables.ImagePath, number, i );
                    Bitmap bitmap = new Bitmap(fileName);
                    bitmaps.Add(bitmap);
                }
            }
            return bitmaps;
        }

        public List<Layer> GetInputLayers(List<Bitmap> bitmaps)
        {
            return
                bitmaps
                    .Select(bitmap => LayerBitmapHelper.CreateLayerFromBitmap(bitmap))
                    .ToList();
        }

        private void MainForm_FormClosing(object sender, FormClosingEventArgs e)
        {
            Stopping = true;
            AskToSaveUnchangedChanges();
        }

        private void AskToSaveUnchangedChanges()
        {
            if (_unsavedChanges == true)
            {

                DialogResult result =
                    MessageBox.Show(
                    this,
                        "You have unsaved changes, would you like to save the DeepBeliefNetwork?",
                        "SharpRBM",
                        MessageBoxButtons.YesNo);

                if (result == DialogResult.Yes)
                {
                    SaveDeepBeliefNetwork();
                }
            }
        }

        private void stopButton_Click(object sender, EventArgs e)
        {
            Stopping = true;
        }

        private void Train(DeepBeliefNetworkTrainer trainer)
        {
            Stopping = false;
            ClearBoxes();
            _unsavedChanges = true;

            int generation = 0;
            while (Stopping == false)
            {
                Stopwatch stopwatch = Stopwatch.StartNew();
                TrainingError error = trainer.Train();
                label1.Text =
                    string.Format(
                    "Gen {0} ({4:0.00} s): ReconstructionError={1:0.00}, DetectorError={2:0.00}, LearningRate={3:0.0000}",
                    generation,
                    error.ReconstructionError,
                    error.FeatureDetectorError,
                    trainer.TrainingWeights.AdjustedLearningRate,
                    stopwatch.ElapsedMilliseconds / 1000.0);

                Application.DoEvents();
                ShowReconstructed(trainer);
                ShowFeatureDetectors(trainer);
                Application.DoEvents();

                if (Stopping)
                {
                    break;
                }
                generation++;
            }

            DocumentDeepBeliefNetwork();
        }

        private void startButton_Click(object sender, EventArgs e)
        {
            int layerId = 0;

            if (trainFirstLayerRadioButton.Checked)
            {
                layerId = 0;
            }
            else if (trainSecondLayerRadioButton.Checked)
            {
                layerId = 1;
            }
            else if (trainThirdLayerRadioButton.Checked)
            {
                layerId = 2;
            }

            AddLog("Preparing to train layer {0}...", layerId);
            DeepBeliefNetworkTrainer trainer =
                new DeepBeliefNetworkTrainer(DeepBeliefNetwork, DeepBeliefNetwork.LayerWeights[layerId], _inputLayers);

            AddLog("Training...");
            Train(trainer);


            AddLog("Training ended");
            AddLog("");
            DocumentDeepBeliefNetwork();
        }

        private void loadButton_Click(object sender, EventArgs e)
        {
            LoadDeepBeliefNetwork();
        }

        private void saveButton_Click(object sender, EventArgs e)
        {
            SaveDeepBeliefNetwork();
        }

        private void SaveDeepBeliefNetwork()
        {
            SerializationHelper.SaveToFile(DeepBeliefNetwork, DeepBeliefNetworkFile);
            _unsavedChanges = false;
        }

        private void clearSelectedLayerButton_Click(object sender, EventArgs e)
        {
            int layerId = 0;
            if (trainFirstLayerRadioButton.Checked)
            {
                layerId = 0;
            }
            else if (trainSecondLayerRadioButton.Checked)
            {
                layerId = 1;
            }
            else if (trainThirdLayerRadioButton.Checked)
            {
                layerId = 2;
            }
            DeepBeliefNetwork.LayerWeights[layerId].CreateRandomWeights();
        }
    }
}
